import numpy
from helper import *

class ParticleSystemConstraints:
    
    def __init__(self, y, x, constraint_radius, constraint_variance, resolution):
        self._constraint_error = numpy.empty((y, x), dtype = numpy.float)
        self._constraint_error[:,:] = 1000000
        self._constraint_window = influence_window(constraint_radius)
        self._constraint_distances = dict()
        for c in self._constraint_window:
            self._constraint_distances[c] = numpy.sqrt(c[0]*c[0] + c[1]*c[1]) * resolution
        
        self._constraint_variance = constraint_variance
    
    def DistanceConstraints(self, points):
        y = points.shape[0]
        x = points.shape[1]
        
        #new point positions
        tmp = numpy.array(points, copy=True)
        
        is_in_range = lambda (i,j): i >= 0 and i < y and j >= 0 and j < x
        norm = lambda x: numpy.sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2])
        
        #calculate new point positions
        for i in range(y):
            for j in range(x):
                v0 = points[i, j]
                for p in self._constraint_window:
                    d = self._constraint_distances[p]
                    k = gauss_2d(p, self._constraint_variance)
                    p = (p[0]+i, p[1]+j)
                    if is_in_range(p):
                        v1 = points[p]
                        v = v1 - v0
                        l = norm(v)
                        n = v / l 
                        dv = (l - d) * n
                        #points[i, j] += k * dv / d    ???
                        #points[i, j] += dv / d        ???
                        tmp[i, j] += k * dv 
        
        points = tmp
        
        for i in range(y):
            for j in range(x):
                v0 = points[i, j]
                self._constraint_error[i, j] = 0
                nbr_constraints = 0
                for p in self._constraint_window:
                    d = self._constraint_distances[p]
                    p = (p[0]+i, p[1]+j)
                    if is_in_range(p):
                        nbr_constraints += 1
                        v1 = points[p]
                        v = v1 - v0
                        l = norm(v)
                        self._constraint_error[i, j] +=  numpy.abs(l - d)
                if nbr_constraints != 0:
                    self._constraint_error[i, j] /= nbr_constraints
                else:
                    self._constraint_error[i, j] = 0
        
        return points
        
                    
    def ConstraintsError(self):
        var = numpy.var(self._constraint_error)
        mean = numpy.mean(self._constraint_error)
        return var / mean + mean 